{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#https://github.com/znxlwm/pytorch-generative-model-collections\n",
    "import utils, torch, time, os, pickle, itertools\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "import torch.utils.data as Utils\n",
    "from skimage import io\n",
    "from PIL import Image\n",
    "import os\n",
    "from matplotlib import pyplot as plt\n",
    "% matplotlib inline\n",
    "\n",
    "# % mogrify -resize 128x128 *.jpg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if not os.path.exists('models'):\n",
    "    os.makedirs('models')\n",
    "\n",
    "# --result_dir\n",
    "if not os.path.exists('results'):\n",
    "    os.makedirs('results')\n",
    "\n",
    "# --result_dir\n",
    "if not os.path.exists('logs'):\n",
    "    os.makedirs('logs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_image(file_name):\n",
    "    data = io.imread(file_name, as_grey=False)\n",
    "    data = np.array(data, order='F')\n",
    "    data = torch.from_numpy(data).transpose(0,2)\n",
    "#     data = np.expand_dims(data,axis=0)\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = load_image(\"Pokemon/abra.jpg\")\n",
    "# a = np.expand_dims(a,axis=4)\n",
    "print(a.shape)\n",
    "a = a.transpose(0,2)\n",
    "print(a.shape)\n",
    "\n",
    "plt.gray()\n",
    "plt.imsave('image.png', a)\n",
    "plt.imshow(a)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class VoiceSets(Utils.Dataset): # Main dataset class for dataloader\n",
    "    def __init__(self, inputs, target): # takes two input, the input and output Tensors\n",
    "        self.inputwave = inputs\n",
    "        self.targetwave = target\n",
    "        self.len = len(self.inputwave) # number of input, NOT LENGTH OF INPUT --> LENGTH OF INPUT = 4000 in our case\n",
    "        \n",
    "    def __len__(self): # must be written for dataset module\n",
    "        return self.len\n",
    "    \n",
    "    def __getitem__(self,index): # must be written for dataset module\n",
    "        return self.inputwave[index], self.targetwave[index] # call for input, target on Training block"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# STANDARD GAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class generator(nn.Module):\n",
    "    def __init__(self, dataset = 'pokemon'):\n",
    "        super(generator, self).__init__()\n",
    "        if dataset == 'mnist' or dataset == 'fashion-mnist':\n",
    "            self.input_height = 28\n",
    "            self.input_width = 28\n",
    "            self.input_dim = 62\n",
    "            self.output_dim = 1\n",
    "        elif dataset == 'celebA':\n",
    "            self.input_height = 64\n",
    "            self.input_width = 64\n",
    "            self.input_dim = 62\n",
    "            self.output_dim = 3\n",
    "        elif dataset == 'pokemon':\n",
    "            self.input_height = 64\n",
    "            self.input_width = 64\n",
    "            self.input_dim = 64\n",
    "            self.output_dim = 3\n",
    "\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(self.input_dim, 1024),\n",
    "            nn.BatchNorm1d(1024),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(1024, 128 * (self.input_height // 4) * (self.input_width // 4)),\n",
    "            nn.BatchNorm1d(128 * (self.input_height // 4) * (self.input_width // 4)),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.deconv = nn.Sequential(\n",
    "            nn.ConvTranspose2d(128, 64, 4, 2, 1),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(64, self.output_dim, 4, 2, 1),\n",
    "            nn.Sigmoid(),\n",
    "        )\n",
    "        utils.initialize_weights(self)\n",
    "\n",
    "    def forward(self, input):\n",
    "        x = self.fc(input)\n",
    "        x = x.view(-1, 128, (self.input_height // 4), (self.input_width // 4))\n",
    "        x = self.deconv(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class discriminator(nn.Module):\n",
    "    def __init__(self, dataset = 'pokemon'):\n",
    "        super(discriminator, self).__init__()\n",
    "        if dataset == 'mnist' or dataset == 'fashion-mnist':\n",
    "            self.input_height = 28\n",
    "            self.input_width = 28\n",
    "            self.input_dim = 1\n",
    "            self.output_dim = 1\n",
    "        elif dataset == 'celebA':\n",
    "            self.input_height = 64\n",
    "            self.input_width = 64\n",
    "            self.input_dim = 3\n",
    "            self.output_dim = 1\n",
    "        elif dataset == 'pokemon':\n",
    "            self.input_height = 64\n",
    "            self.input_width = 64\n",
    "            self.input_dim = 3\n",
    "            self.output_dim = 1\n",
    "            \n",
    "        self.conv = nn.Sequential(\n",
    "            nn.Conv2d(self.input_dim, 64, 4, 2, 1),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Conv2d(64, 128, 4, 2, 1),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.LeakyReLU(0.2),\n",
    "        )\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(128 * (self.input_height // 4) * (self.input_width // 4), 1024),\n",
    "            nn.BatchNorm1d(1024),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Linear(1024, self.output_dim),\n",
    "            nn.Sigmoid(),\n",
    "        )\n",
    "        utils.initialize_weights(self)\n",
    "\n",
    "    def forward(self, input):\n",
    "        x = self.conv(input)\n",
    "        x = x.view(-1, 128 * (self.input_height // 4) * (self.input_width // 4))\n",
    "        x = self.fc(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GAN(object):\n",
    "    def __init__(self):\n",
    "        # parameters\n",
    "        self.epoch = 500\n",
    "        self.sample_num = 32\n",
    "        self.batch_size = 32\n",
    "        self.save_dir = 'models'\n",
    "        self.result_dir = 'results'\n",
    "        self.dataset = 'pokemon'\n",
    "        self.log_dir = 'logs'\n",
    "        self.gpu_mode = True\n",
    "        self.model_name = 'GAN'\n",
    "        self.lrG = 0.0001\n",
    "        self.lrD = 0.0001\n",
    "        self.beta1 = 0.4\n",
    "        self.beta2 = 0.999\n",
    "        \n",
    "\n",
    "        # networks initmnist\n",
    "        self.G = generator(self.dataset)\n",
    "        self.D = discriminator(self.dataset)\n",
    "        self.G_optimizer = optim.Adam(self.G.parameters(), lr=self.lrG, betas=(self.beta1, self.beta2))\n",
    "        self.D_optimizer = optim.Adam(self.D.parameters(), lr=self.lrD, betas=(self.beta1, self.beta2))\n",
    "\n",
    "        if self.gpu_mode:\n",
    "            self.G.cuda()\n",
    "            self.D.cuda()\n",
    "            self.BCE_loss = nn.BCELoss().cuda()\n",
    "        else:\n",
    "            self.BCE_loss = nn.BCELoss()\n",
    "\n",
    "        # load dataset\n",
    "        if self.dataset == 'mnist':\n",
    "            self.data_loader = DataLoader(datasets.MNIST('data/mnist', train=True, download=True,\n",
    "                                                                          transform=transforms.Compose(\n",
    "                                                                              [transforms.ToTensor()])),\n",
    "                                                           batch_size=self.batch_size, shuffle=True)\n",
    "        elif self.dataset == 'fashion-mnist':\n",
    "            self.data_loader = DataLoader(\n",
    "                datasets.FashionMNIST('data/fashion-mnist', train=True, download=True, transform=transforms.Compose(\n",
    "                    [transforms.ToTensor()])),\n",
    "                batch_size=self.batch_size, shuffle=True)\n",
    "        elif self.dataset == 'celebA':\n",
    "            self.data_loader = utils.load_celebA('data/celebA', transform=transforms.Compose(\n",
    "                [transforms.CenterCrop(160), transforms.Scale(64), transforms.ToTensor()]), batch_size=self.batch_size,\n",
    "                                                 shuffle=True)\n",
    "        elif self.dataset == 'pokemon':\n",
    "            fs = !ls Pokemon *.jpg\n",
    "            fs = fs[2:]\n",
    "            poke_trn = []\n",
    "            for name in fs:\n",
    "                poke_trn.append(load_image('Pokemon/' + name))  \n",
    "                \n",
    "            fs = !ls FlipPokemon *.jpg\n",
    "            fs = fs[2:]\n",
    "            for name in fs:\n",
    "                poke_trn.append(load_image('FlipPokemon/' + name))\n",
    "            dataset = VoiceSets(poke_trn, poke_trn)\n",
    "            self.data_loader = Utils.DataLoader(dataset=dataset, shuffle = True, batch_size=self.batch_size)\n",
    "            \n",
    "        self.z_dim = 64\n",
    "\n",
    "        # fixed noise\n",
    "        if self.gpu_mode:\n",
    "            self.sample_z_ = Variable(torch.rand((self.batch_size, self.z_dim)).cuda(), volatile=True)\n",
    "        else:\n",
    "            self.sample_z_ = Variable(torch.rand((self.batch_size, self.z_dim)), volatile=True)\n",
    "    \n",
    "    def train(self):\n",
    "        self.train_hist = {}\n",
    "        self.train_hist['D_loss'] = []\n",
    "        self.train_hist['G_loss'] = []\n",
    "        self.train_hist['per_epoch_time'] = []\n",
    "        self.train_hist['total_time'] = []\n",
    "\n",
    "        if self.gpu_mode:\n",
    "            self.y_real_, self.y_fake_ = Variable(torch.ones(self.batch_size, 1).cuda()), Variable(torch.zeros(self.batch_size, 1).cuda())\n",
    "        else:\n",
    "            self.y_real_, self.y_fake_ = Variable(torch.ones(self.batch_size, 1)), Variable(torch.zeros(self.batch_size, 1))\n",
    "\n",
    "        self.D.train()\n",
    "        print('training start!!')\n",
    "        start_time = time.time()\n",
    "        for epoch in range(self.epoch):\n",
    "            self.G.train()\n",
    "            epoch_start_time = time.time()\n",
    "            for iter, (x_,_) in enumerate(self.data_loader):\n",
    "                if iter == self.data_loader.dataset.__len__() // self.batch_size:\n",
    "                    break\n",
    "                    \n",
    "                x_ = x_.type(torch.FloatTensor)\n",
    "                \n",
    "                z_ = torch.rand((self.batch_size, self.z_dim))\n",
    "               \n",
    "                if self.gpu_mode:\n",
    "                    x_, z_ = Variable(x_.cuda()), Variable(z_.cuda())\n",
    "                else:\n",
    "                    x_, z_ = Variable(x_), Variable(z_)\n",
    "    \n",
    "                # update D network\n",
    "                self.D_optimizer.zero_grad()\n",
    "                \n",
    "                D_real = self.D(x_)\n",
    "                D_real_loss = self.BCE_loss(D_real, self.y_real_)\n",
    "\n",
    "                G_ = self.G(z_)\n",
    "\n",
    "                D_fake = self.D(G_)\n",
    "                D_fake_loss = self.BCE_loss(D_fake, self.y_fake_)\n",
    "\n",
    "                D_loss = D_real_loss + D_fake_loss\n",
    "                self.train_hist['D_loss'].append(D_loss.data[0])\n",
    "\n",
    "                D_loss.backward()\n",
    "                self.D_optimizer.step()\n",
    "\n",
    "                # update G network\n",
    "                self.G_optimizer.zero_grad()\n",
    "\n",
    "                G_ = self.G(z_)\n",
    "                D_fake = self.D(G_)\n",
    "                G_loss = self.BCE_loss(D_fake, self.y_real_)\n",
    "                self.train_hist['G_loss'].append(G_loss.data[0])\n",
    "\n",
    "                G_loss.backward()\n",
    "                self.G_optimizer.step()\n",
    "\n",
    "                if ((iter + 1) % 100) == 0:\n",
    "                    print(\"Epoch: [%2d] [%4d/%4d] D_loss: %.8f, G_loss: %.8f\" %\n",
    "                          ((epoch + 1), (iter + 1), self.data_loader.dataset.__len__() // self.batch_size, D_loss.data[0], G_loss.data[0]))\n",
    "\n",
    "            self.train_hist['per_epoch_time'].append(time.time() - epoch_start_time)\n",
    "            if ((epoch + 1)%25 == 0) or (epoch == 0):\n",
    "                self.visualize_results((epoch+1)/25)\n",
    "\n",
    "        self.train_hist['total_time'].append(time.time() - start_time)\n",
    "        print(\"Avg one epoch time: %.2f, total %d epochs time: %.2f\" % (np.mean(self.train_hist['per_epoch_time']),\n",
    "              self.epoch, self.train_hist['total_time'][0]))\n",
    "        print(\"Training finish!... save training results\")\n",
    "\n",
    "        self.save()\n",
    "        utils.generate_animation(self.result_dir + '/' + self.dataset + '/' + self.model_name + '/' + self.model_name,\n",
    "                                 ((self.epoch+1)//25))\n",
    "        utils.loss_plot(self.train_hist, os.path.join(self.save_dir, self.dataset, self.model_name), self.model_name)\n",
    "\n",
    "    def visualize_results(self, epoch, fix=False):\n",
    "        self.G.eval()\n",
    "\n",
    "        if not os.path.exists(self.result_dir + '/' + self.dataset + '/' + self.model_name):\n",
    "            os.makedirs(self.result_dir + '/' + self.dataset + '/' + self.model_name)\n",
    "\n",
    "        tot_num_samples = min(self.sample_num, self.batch_size)\n",
    "        image_frame_dim = int(np.floor(np.sqrt(tot_num_samples)))\n",
    "\n",
    "        if fix:\n",
    "            \"\"\" fixed noise \"\"\"\n",
    "            samples = self.G(self.sample_z_)\n",
    "        else:\n",
    "            \"\"\" random noise \"\"\"\n",
    "            if self.gpu_mode:\n",
    "                sample_z_ = Variable(torch.rand((self.batch_size, self.z_dim)).cuda(), volatile=True)\n",
    "            else:\n",
    "                sample_z_ = Variable(torch.rand((self.batch_size, self.z_dim)), volatile=True)\n",
    "            samples = self.G(sample_z_)\n",
    "        if self.gpu_mode:\n",
    "            samples = samples.cpu().data.numpy().transpose(0, 3, 2, 1)\n",
    "        else:\n",
    "            samples = samples.data.numpy().transpose(0, 2, 3, 1)\n",
    "        utils.save_images(samples[:image_frame_dim * image_frame_dim, :, :, :], [image_frame_dim, image_frame_dim],\n",
    "                          self.result_dir + '/' + self.dataset + '/' + self.model_name + '/' + self.model_name + '_epoch%03d' % epoch + '.png')\n",
    "\n",
    "    def save(self):\n",
    "        save_dir = os.path.join(self.save_dir, self.dataset, self.model_name)\n",
    "\n",
    "        if not os.path.exists(save_dir):\n",
    "            os.makedirs(save_dir)\n",
    "\n",
    "        torch.save(self.G.state_dict(), os.path.join(save_dir, self.model_name + '_G.pkl'))\n",
    "        torch.save(self.D.state_dict(), os.path.join(save_dir, self.model_name + '_D.pkl'))\n",
    "\n",
    "        with open(os.path.join(save_dir, self.model_name + '_history.pkl'), 'wb') as f:\n",
    "            pickle.dump(self.train_hist, f)\n",
    "\n",
    "    def load(self):\n",
    "        save_dir = os.path.join(self.save_dir, self.dataset, self.model_name)\n",
    "\n",
    "        self.G.load_state_dict(torch.load(os.path.join(save_dir, self.model_name + '_G.pkl')))\n",
    "        self.D.load_state_dict(torch.load(os.path.join(save_dir, self.model_name + '_D.pkl')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "gan = GAN()\n",
    "\n",
    "gan.train()\n",
    "print(\" [*] Training finished!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "utils.generate_animation('results/pokemon/GAN/GAN', 15)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:wavenet]",
   "language": "python",
   "name": "conda-env-wavenet-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
